import socket
import struct
import threading
import zlib
import time
import pyautogui
import pyaudio
from io import BytesIO
from PIL import Image


class RemoteServer:
    def __init__(self, host='0.0.0.0', port=65432):
        self.server_socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        self.server_socket.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
        self.server_socket.bind((host, port))
        self.server_socket.listen(5)
        self.running = True
        self.target_fps = 30
        self.quality = 70
        self.frame_interval = 1.0 / self.target_fps

        # 音频配置
        self.audio_format = pyaudio.paInt16
        self.audio_channels = 1
        self.audio_rate = 44100
        self.audio_chunk = 1024
        self.audio = pyaudio.PyAudio()

        print(f"远程桌面服务端启动在 {host}:{port}")

    def capture_screen(self):
        """捕获屏幕并压缩"""
        img = pyautogui.screenshot()
        with BytesIO() as buf:
            img.save(buf, format='JPEG', quality=self.quality)
            return zlib.compress(buf.getvalue())

    def capture_audio(self):
        """捕获音频并压缩"""
        stream = self.audio.open(
            format=self.audio_format,
            channels=self.audio_channels,
            rate=self.audio_rate,
            input=True,
            frames_per_buffer=self.audio_chunk
        )
        data = stream.read(self.audio_chunk)
        stream.stop_stream()
        stream.close()
        return zlib.compress(data)

    def handle_input(self, client_socket):
        """处理客户端输入（鼠标/键盘）"""
        try:
            while self.running:
                input_type = client_socket.recv(1)
                if not input_type:
                    break

                if input_type == b'K':
                    # 处理键盘输入（4字节键码）
                    key_data = client_socket.recv(4)
                    key_code = struct.unpack('!I', key_data)[0]
                    pyautogui.press(chr(key_code))

                elif input_type == b'M':
                    # 处理鼠标移动（8字节坐标：x(4字节)+y(4字节)）
                    pos_data = client_socket.recv(8)
                    x, y = struct.unpack('!II', pos_data)
                    pyautogui.moveTo(x, y)

                elif input_type == b'C':
                    # 处理鼠标点击（1字节按钮：L/R）
                    btn_data = client_socket.recv(1)
                    if btn_data == b'L':
                        pyautogui.click(button='left')
                    elif btn_data == b'R':
                        pyautogui.click(button='right')
        except Exception as e:
            print(f"输入处理错误: {e}")

    def handle_client(self, client_socket, addr):
        """处理客户端连接（发送音视频+接收输入）"""
        print(f"新客户端连接: {addr}")
        # 启动输入处理线程
        input_thread = threading.Thread(
            target=self.handle_input,
            args=(client_socket,),
            daemon=True
        )
        input_thread.start()

        try:
            last_frame_time = 0
            last_audio_time = 0
            audio_interval = 0.1  # 100ms音频采样间隔

            while self.running:
                current_time = time.time()

                # 发送视频帧（控制帧率）
                if current_time - last_frame_time >= self.frame_interval:
                    frame_data = self.capture_screen()
                    try:
                        # 协议：[类型(V)+长度(4字节)+数据]
                        client_socket.sendall(b'V' + struct.pack('!I', len(frame_data)) + frame_data)
                        last_frame_time = current_time
                    except:
                        break

                # 发送音频数据
                if current_time - last_audio_time >= audio_interval:
                    audio_data = self.capture_audio()
                    try:
                        # 协议：[类型(A)+长度(4字节)+数据]
                        client_socket.sendall(b'A' + struct.pack('!I', len(audio_data)) + audio_data)
                        last_audio_time = current_time
                    except:
                        break

        except Exception as e:
            print(f"客户端 {addr} 处理错误: {e}")
        finally:
            client_socket.close()
            print(f"客户端 {addr} 断开连接")

    def start(self):
        """启动服务端"""
        try:
            while self.running:
                client_socket, addr = self.server_socket.accept()
                # 为每个客户端创建独立线程
                threading.Thread(
                    target=self.handle_client,
                    args=(client_socket, addr),
                    daemon=True
                ).start()
        except KeyboardInterrupt:
            print("正在关闭服务端...")
        finally:
            self.audio.terminate()
            self.server_socket.close()

    def stop(self):
        """停止服务端"""
        self.running = False


if __name__ == "__main__":
    server = RemoteServer()
    server.start()